#pragma once
#include "psram_manager.h"

Memory::Memory() : head(nullptr) {}

void Memory::init() {
  this->maxSramMemory = heap_caps_get_total_size(MALLOC_CAP_INTERNAL);
  this->maxFlashMemory = ESP.getFlashChipSize();

  if (psramFound())
  {
    LOG("PSRAM is available!\n");

    this->isPsramFound = true;
    this->maxPsramMemory = ESP.getPsramSize();
  }
  else
  {
    LOG("PSRAM not available!\n");
    this->isPsramFound = false;
  }

  LOG("Max SRAM Memory: ", this->maxSramMemory / 1024, " KB\nMax Flash Memory: ", this->maxFlashMemory / 1024, " KB\n");
  if (this->isPsramFound) {
    LOG("Max PSRAM Memory: ", this->maxPsramMemory / 1024, " KB\n");
  }
  LOG("Used SRAM Memory: ", (this->maxSramMemory -  heap_caps_get_free_size(MALLOC_CAP_INTERNAL))  / 1024, " KB\n");
  if (this->isPsramFound) {
    LOG("Used PSRAM Memory: ",  (this->maxPsramMemory - ESP.getFreePsram()) / 1024, " KB\n");
  }
}

size_t Memory::getUsedSramMemory()
{
  return this->maxSramMemory -  heap_caps_get_free_size(MALLOC_CAP_INTERNAL); 
}

size_t Memory::getUsedPsramMemory()
{
  if (!this->isPsramFound)
  {
    LOG("PSRAM not available!\n");
  }
  return this->maxPsramMemory - ESP.getFreePsram();
}

size_t Memory::getMaxSramMemory() {
  return this->maxSramMemory;
}

size_t Memory::getMaxPsramMemory() {
  if (!this->isPsramFound)
  {
    LOG("PSRAM not available!\n");
  }
  return this->maxPsramMemory;
}

void *Memory::allocate(size_t size, uint8_t mode)
{
  void *ptr = nullptr;

  if (mode == SRAM)
  {
    if (heap_caps_get_free_size(MALLOC_CAP_INTERNAL) < size)
    {
      LOG("Memory allocation failed: not enough SRAM memory, need ", size / 1024, " KB, but only ",  heap_caps_get_free_size(MALLOC_CAP_INTERNAL) / 1024, " KB available\n");
      return nullptr;
    }
    else
    {
      ptr = malloc(size);
    }
  }
  else if (mode == PSRAM)
  {
    if (this->isPsramFound)
    {
      if (ESP.getFreePsram() < size)
      {
        LOG("Memory allocation failed: not enough PSRAM memory, need ", size / 1024, " KB, but only ", ESP.getFreePsram() / 1024, " KB available\n");
        return nullptr;
      }
      else
      {
        ptr = heap_caps_malloc(size, MALLOC_CAP_SPIRAM);
      }
    }
    else
    {
      LOG("PSRAM not available!\n");
      return nullptr;
    }
  }
  if (ptr) {
    MemoryBlock* newBlock = (MemoryBlock*)malloc(sizeof(MemoryBlock));
    newBlock->ptr = ptr;
    newBlock->blockSize = size;
    newBlock->mode = mode;
    newBlock->next = this->head;
    this->head = newBlock;
    String currentPtrAddress = String((uint32_t)newBlock->ptr, HEX);  // 将地址转换为十六进制字符串
    if (newBlock->blockSize > 10240) {
      LOG("Add Memory Block: Adress:", currentPtrAddress, "  Size:", newBlock->blockSize / 1024, "KB   Mode:", newBlock->mode, "\n");
      String currentHeadAddress = String((uint32_t)this->head->ptr, HEX);  // 将地址转换为十六进制字符串
      LOG("Head Memory Block: Adress:", currentHeadAddress, "\n");
      if (mode == SRAM) {

        LOG("SRAM Memory Allocated: ", size / 1024, " KB, and have ",  heap_caps_get_free_size(MALLOC_CAP_INTERNAL) / 1024, " KB left\n");
      } else if (mode == PSRAM) {
        LOG("PSRAM Memory Allocated: ", size / 1024, " KB, and have ", ESP.getFreePsram() / 1024, " KB left\n");
      } 
    }
  } else {
    if (mode == SRAM) {
      LOG("SRAM Memory allocation failed!\n");
    } else if (mode == PSRAM) {
      LOG("PSRAM Memory allocation failed!\n");
    }
  }
  return ptr;
}

void Memory::deallocate(void *ptr)
{
  if (ptr == nullptr) {
    LOG("Memory deallocation failed: pointer is null\n");
    return;
  }

  // LOG("Memory Deallocate\n");
  MemoryBlock* current = this->head;
  MemoryBlock* prev = nullptr;
  // String headPtrAddress = String((uint32_t)current->ptr, HEX);  // 将地址转换为十六进制字符串
  // LOG("Head Memory Block: Adress:", headPtrAddress, "\n");

  while (current->ptr != nullptr) {
    if (current->ptr == ptr) {
      if (current->blockSize > 10240) {
        LOG("Find Memory Block to Deallocate\n");
        LOG("Memory Block Mode:", current->mode, "\n");
        LOG("Memory Block Size:", current->blockSize / 1024, "KB\n");
      }

      if (prev) {
        prev->next = current->next;
      } else {
        this->head = current->next;
      }

      if (current->mode == SRAM) {
        free(ptr);
        if (current->blockSize > 10240) {
          LOG("SRAM Memory Deallocated: ",  current->blockSize / 1024, " KB, and have ",  heap_caps_get_free_size(MALLOC_CAP_INTERNAL) / 1024, " KB left\n");
        }
      } else if (current->mode == PSRAM) {
        heap_caps_free(ptr);
        LOG("PSRAM Memory Deallocated: ", current->blockSize / 1024, " KB, and have ", ESP.getFreePsram() / 1024, " KB left\n");
      }
      // delete current;
      free(current);
      if (current->blockSize > 10240) {
        LOG("Remove MemoryBlock, Memory Deallocate finished.\n");
      }
      return;
    } else {
      String currentPtrAddress = String((uint32_t)current->ptr, HEX);  // 将地址转换为十六进制字符串
      // LOG("Current Chain Node's Pointer Address: ", currentPtrAddress.c_str(), "\n");  // 打印
    }
    prev = current;
    current = current->next;
  }
  LOG("Memory Deallocate failed: pointer not found\n");
}

MemoryBlock* Memory::getHead()
{
  return this->head;
}

void Memory::clear()
{
  LOG("Clearing all allocated memory blocks\n");

  MemoryBlock* current = this->head;
  MemoryBlock* next = nullptr;

  while (current) {
      next = current->next;

      // // Log memory block details
      // LOG("Deallocating Memory Block\n");
      // LOG("Memory Block Mode:", current->mode, "\n");
      // LOG("Memory Block Size:", current->blockSize, " bytes\n");

      // Free the memory based on the mode
      if (current->mode == SRAM) {
          free(current->ptr);
          if (current->blockSize > 10240) {
            LOG("SRAM Memory Deallocated: ", current->blockSize / 1024, " KB\n");
          }
      } else if (current->mode == PSRAM) {
          heap_caps_free(current->ptr);
          LOG("PSRAM Memory Deallocated: ", current->blockSize / 1024, " KB\n");
      }

      // Free the MemoryBlock itself
      free(current);
      LOG("MemoryBlock removed\n");

      current = next;  // Move to the next block
  }

  // Reset the head to nullptr
  this->head = nullptr;
  LOG("All memory blocks cleared\n");
}

Memory::~Memory()
{
  clear();
}

Memory memory;

void initMemory()
{
  memory.init();
}

void *pMalloc(size_t n)
{
  if (n > 50 * 1024)
  {
    return memory.allocate(n, PSRAM);
  }
  return memory.allocate(n, SRAM);
}

void Free(void *ptr)
{
  memory.deallocate(ptr);
}

// bytes
size_t getMaxSramMemory()
{
  return memory.getMaxSramMemory();
}

size_t getMaxPsramMemory() {
  return memory.getMaxPsramMemory();
}

size_t getUsedSramMemory()
{
  return memory.getUsedSramMemory();
}

size_t getUsedPsramMemory()
{
  return memory.getUsedPsramMemory();
}

size_t getFreeSramMemory() {
  return memory.getMaxSramMemory() - memory.getUsedSramMemory();
}

size_t getFreePsramMemory() {
  return memory.getMaxPsramMemory() - memory.getUsedPsramMemory();
}

void ClearMemory()
{
  memory.clear();
}

void getHeadAdress(int a) {
  MemoryBlock* current = memory.getHead();
  String headAdress = String((uint32_t)current->ptr, HEX);  // 将地址转换为十六进制字符串
  LOG(a, ":  Head Memory Block: Adress:", headAdress, "\n");
}
void printMemoryInfo(char *msg) {
  LOG(msg, "\nSRAM used: ", getUsedSramMemory() / 1024, " KB, remaining: ", getFreeSramMemory() / 1024, " KB, PSRAM used: ", getUsedPsramMemory() / 1024, " KB, remaining: ", getFreePsramMemory() / 1024, "KB.\n");
}